"""
水库群供水调度优化模型数据处理器
将set_up.py加载的原始JSON数据转换为优化模型所需的数据结构
"""

from typing import Dict, List, Any, Tuple
from loguru import logger

class DataProcessor:
    """数据处理器，将set_up.py加载的原始数据转换为优化模型所需的数据结构"""
    
    def __init__(self, para: Dict[str, Any]):
        """
        初始化数据处理器
        
        Args:
            para: 从set_up.py加载的参数字典，包含config、runoff_data、demand_data、obj_cons_data
        """
        self.para = para
        self.config = para['config']
        self.runoff_data = para['runoff_data']
        self.demand_data = para['demand_data']
        self.obj_cons_data = para.get('obj_cons_data', {})
    
    def process_data(self) -> Dict[str, Any]:
        """
        处理数据并返回优化模型所需的数据结构
        
        Returns:
            包含所有处理后数据的字典
        """
        param = {}
        
        try:
            # 1. 处理径流数据
            param['runoff'] = self._process_runoff()
            logger.info("径流数据处理完成")
            
            # 2. 处理需水数据
            param['demand'] = self._process_demand()
            logger.info("需水数据处理完成")
            
            # 3. 处理供给关系
            param['supply_relationship'] = self._process_supply_relationship()
            logger.info("供给关系处理完成")
            param['reservoir_supply'] = self._process_reservoir_supply()
            logger.info("水库-用水户映射处理完成")
            
            # 4. 处理初始库容
            param['storage_init'] = self._process_storage_init()
            logger.info("初始库容处理完成")
            
            # 4.1 处理汛末库容约束（end_storage_flood_season）
            param['end_storage_flood_season_cons'] = self._process_end_storage_flood_season_cons()
            logger.info("汛末库容约束处理完成（end_storage_flood_season_cons）")
            
            # 4.2 处理调度期末库容约束（end_storage_dispatch_term）
            param['end_storage_dispatch_term_cons'] = self._process_end_storage_dispatch_term_cons()
            logger.info("调度期末库容约束处理完成（end_storage_dispatch_term_cons）")
            
            # 5. 处理库容上限
            param['storage_up'] = self._process_storage_up()
            logger.info("库容上限处理完成")
            
            # 6. 处理库容下限
            param['storage_down'] = self._process_storage_down()
            logger.info("库容下限处理完成")

            # 6.1 处理弃水上限约束（spillflow_up）
            param['spillflow_up'] = self._process_spillflow_up()
            logger.info("弃水约束处理完成（spillflow_up）")
            
            # 6.2 处理弃水下限约束（spillflow_down）
            param['spillflow_down'] = self._process_spillflow_down()
            logger.info("弃水下限约束处理完成（spillflow_down）")
            
            # 7. 处理供水关系类型
            param['supply_relationship_type'] = self._process_supply_relationship_type()
            logger.info("供水关系类型处理完成")
            
            # 8. 处理静态信息
            param['static'] = self._process_static()
            logger.info("静态信息处理完成")
            
            # 9. 处理时间序列索引
            param['time'] = self._process_time_sets()
            logger.info("时间序列索引处理完成")
            
            # 10. 处理水库工程信息
            param['reservoir'] = self._process_engineering_hydro()
            logger.info("水库工程信息处理完成")
            
            # 11. 处理用水户信息
            param['water_user'] = self._process_engineering_water_user()
            logger.info("用水户信息处理完成")
            
            # 12. 处理外部引入水信息
            param['external_water_sources'] = self._process_external_water_sources()
            logger.info("外部引入水信息处理完成")
            
            # 13. 处理外部水源目标水库信息
            param['imported_reservoir'] = self._process_imported_reservoir()
            logger.info("外部水源目标水库信息处理完成")
            
            # 14. 处理外部水源与目标水库关系
            param['import_relationship'] = self._process_import_relationship()
            logger.info("外部水源与目标水库关系处理完成")
            param['import_relationship_source'] = self._process_import_relationship_source()
            logger.info("外部水源与目标水库关系处理完成")
            
            # 15. 处理损失率
            param['loss_rate'] = self._process_loss_rate()
            logger.info("损失率处理完成")
            
            # 16. 处理输水能力
            param['trans_capacity'] = self._process_trans_capacity()
            logger.info("输水能力处理完成")
            
            # 17. 处理水库间连接关系
            param['reservoir_connections'] = self._process_reservoir_connections()
            logger.info("水库间连接关系处理完成")

            # 18. 处理外部水源引入上限（以目标水库为索引）
            param['import_export_up'] = self._process_import_export_up()
            logger.info("外部水源引入上限(import_export_up)处理完成")
            
            # 18.1 处理外部水源引入下限（以目标水库为索引）
            param['import_export_down'] = self._process_import_export_down()
            logger.info("外部水源引入下限(import_export_down)处理完成")

            # 19. 处理目标函数配置（objective）：支持未来多目标
            objectives = self._process_objectives_config()
            param['objectives'] = objectives
            logger.info("目标函数配置处理完成 (objectives)")
            
        except Exception as e:
            logger.error(f"处理数据时出错: {e}")
            logger.error(f"错误发生在处理: {param.keys()}")
            raise
        
        # 将原始配置参数也包含在返回的param中
        # 这样output_filename等配置信息就能传递到model.param中
        if 'output_filename' in self.para:
            param['output_filename'] = self.para['output_filename']
        if 'output_folder' in self.para:
            param['output_folder'] = self.para['output_folder']
        if 'config_data' in self.para:
            param['config_data'] = self.para['config_data']
        if 'input_filepath' in self.para:
            param['input_filepath'] = self.para['input_filepath']
        if 'command_line_args' in self.para:
            param['command_line_args'] = self.para['command_line_args']
        
        # 添加solver配置传递
        try:
            solver_name = self.para['config']['general_parameters']['solver']
            param['solver'] = {'name': solver_name}
        except (KeyError, TypeError):
            # 如果配置中没有solver字段或配置结构不正确，使用默认值
            param['solver'] = {'name': 'highs'}
        
        return param

    def _process_objectives_config(self):
        """Parse objective configuration from obj_cons.json.

        Returns
        -------
        List[Dict[str, Any]]
            A list of objective entries. Each entry includes the field
            `sense` and optional `name` and `scope`. When `scope` is 'unit',
            the entry must include `unit_stcd` and `unit_stnm`. Downstream
            logic primarily uses `unit_stcd` to reference the target reservoir.
        """
        # 读取配置（顶层 objective）
        objective_cfg = self.obj_cons_data.get('objective', None) if isinstance(self.obj_cons_data, dict) else None

        objectives = []

        for item in objective_cfg:
            # 向后兼容：优先使用 importance，如果不存在则尝试使用 weight
            importance_value = item.get('importance')
            if importance_value is None:
                # 兼容旧配置的 weight 字段
                importance_value = item.get('weight')
            
            # 允许任意对象结构；透传字段
            obj_entry: Dict[str, Any] = {
                'name': item.get('name'),
                'sense': item.get('sense'),
                'scope': item.get('scope'),
                'importance': importance_value  # 可以是语气词或数值
            }

            # unit 级别需要额外的字段
            if item.get('scope') == 'unit':
                obj_entry['unit_stcd'] = item.get('unit_stcd')
                obj_entry['unit_stnm'] = item.get('unit_stnm')

            objectives.append(obj_entry)

        # 将模糊语气转换为归一化权重
        self._convert_importance_to_weights(objectives)

        return objectives
    
    def _convert_importance_to_weights(self, objectives: List[Dict[str, Any]]) -> None:
        """将模糊语气重要性转换为归一化权重。
        
        模糊语气分为5个等级：
        - 非常重要: 5分
        - 重要: 4分
        - 一般: 3分
        - 不重要: 2分
        - 非常不重要: 1分
        
        根据各目标的分数，计算归一化权重，确保权重总和为1。
        
        Parameters
        ----------
        objectives : List[Dict[str, Any]]
            目标函数配置列表，会在原地修改，添加 'weight' 字段
        """
        # 定义语气词到分数的映射
        importance_scores = {
            '非常重要': 5,
            '重要': 4,
            '一般': 3,
            '不重要': 2,
            '非常不重要': 1
        }
        
        # 计算每个目标的分数
        scores = []
        for obj in objectives:
            importance = obj.get('importance', '一般')
            
            # 兼容旧配置：如果已经是数值，直接使用
            if isinstance(importance, (int, float)):
                score = float(importance)
            else:
                # 去除空白字符并查找匹配
                importance_str = str(importance).strip()
                score = importance_scores.get(importance_str, 3)  # 默认为"一般"(3分)
                
                if importance_str not in importance_scores:
                    logger.warning(
                        f"未识别的重要性语气 '{importance_str}'，使用默认值'一般'(3分)"
                    )
            
            scores.append(score)
        
        # 计算总分
        total_score = sum(scores)
        
        if total_score == 0:
            # 如果总分为0，平均分配
            logger.warning("所有目标的重要性分数总和为0，采用平均分配权重")
            weight = 1.0 / len(objectives) if objectives else 0
            for obj in objectives:
                obj['weight'] = weight
        else:
            # 归一化为权重
            for obj, score in zip(objectives, scores):
                obj['weight'] = score / total_score
        
        # 记录权重分配结果
        logger.info("目标函数权重分配结果：")
        for obj in objectives:
            importance = obj.get('importance', '未指定')
            weight = obj.get('weight', 0)
            name = obj.get('name', '未命名')
            scope = obj.get('scope', 'system')
            unit_stcd = obj.get('unit_stcd', '')
            
            if scope == 'unit' and unit_stcd:
                logger.info(f"  - {name}({unit_stcd}): 重要性='{importance}', 权重={weight:.4f}")
            else:
                logger.info(f"  - {name}(系统级): 重要性='{importance}', 权重={weight:.4f}")
    
    def _process_runoff(self) -> Dict[Tuple[str, str], float]:
        """
        处理径流数据
        
        Returns:
            字典，key为(水库名称, 时间)的元组，value为径流量
        """
        runoff = {}
        
        # 检查runoff_data的结构
        if 'reservoirs' in self.runoff_data:
            # 如果runoff_data包含reservoirs键，则处理列表格式
            reservoirs_list = self.runoff_data['reservoirs']
            for reservoir in reservoirs_list:
                reservoir_stcd = reservoir.get('stcd')
                reservoir_name = reservoir.get('stnm')
                
                if not reservoir_stcd or not reservoir_name:
                    logger.warning(f"水库数据缺少必要字段: {reservoir}")
                    continue
                
                # 检查配置中是否有该水库
                reservoirs_config = self.config.get('reservoirs', [])
                reservoir_info = None
                if isinstance(reservoirs_config, dict):
                    reservoir_info = reservoirs_config.get(reservoir_stcd)
                else:
                    for res in reservoirs_config:
                        if res.get('stcd') == reservoir_stcd:
                            reservoir_info = res
                            break
                
                if not reservoir_info:
                    logger.warning(f"未找到水库 {reservoir_stcd} 的配置信息")
                    continue
                
                # 处理径流数据
                runoff_data = reservoir.get('runoff_data', {})
                for time_key, flow_value in runoff_data.items():
                    # 直接使用time_key作为时间标识，格式为 "YYYY-MM"
                    if isinstance(time_key, str):
                        runoff[(reservoir_stcd, time_key)] = flow_value
        else:
            # 如果runoff_data是字典格式（旧格式兼容）
            for reservoir_stcd, reservoir_runoff in self.runoff_data.items():
                # 获取水库名称
                reservoirs_config = self.config.get('reservoirs', [])
                reservoir_info = None
                if isinstance(reservoirs_config, dict):
                    reservoir_info = reservoirs_config.get(reservoir_stcd)
                else:
                    for res in reservoirs_config:
                        if res.get('stcd') == reservoir_stcd:
                            reservoir_info = res
                            break
                
                if not reservoir_info:
                    logger.warning(f"未找到水库 {reservoir_stcd} 的配置信息")
                    continue
                
                reservoir_stcd = reservoir_info['stcd']
                
                # 处理径流数据
                for time_key, flow_value in reservoir_runoff.items():
                    # 直接使用time_key作为时间标识，格式为 "YYYY-MM"
                    if isinstance(time_key, str):
                        runoff[(reservoir_stcd, time_key)] = flow_value
        
        return runoff
    
    def _process_demand(self) -> Dict[Tuple[str, str], float]:
        """
        处理需水数据
        
        Returns:
            字典，key为(用水户名称, 时间)的元组，value为需水量
        """
        demand = {}
        
        # 检查demand_data的结构
        if 'water_users' in self.demand_data:
            # 如果demand_data包含water_users键，则处理列表格式
            water_users_list = self.demand_data['water_users']
            for user in water_users_list:

                user_stcd = user.get('user_stcd')
                # 优先从demand.json中获取用户名，如果没有则从config中获取
                user_stcd_from_data = user.get('user_stcd')
                
                if not user_stcd:
                    logger.warning(f"用水户数据缺少user_stcd字段: {user}")
                    continue
                
                # 如果demand.json中没有用户名，则从config中获取
                if not user_stcd_from_data:
                    water_users_config = self.config.get('water_users', [])
                    if isinstance(water_users_config, dict):
                        water_users_list = water_users_config.values()
                    else:
                        water_users_list = water_users_config
                    
                    for config_user in water_users_list:
                        if config_user.get('user_stcd') == user_stcd:
                            user_stcd_from_data = config_user.get('user_stcd', user_stcd)
                            break
                    if not user_stcd_from_data:
                        user_stcd_from_data = user_stcd
                        logger.warning(f"未找到用水户 {user_stcd} 的名称，使用原名称")
                
                # 处理需水数据
                demand_data = user.get('demand_data', {})
                for time_key, demand_value in demand_data.items():
                    # 直接使用time_key作为时间标识，格式为 "YYYY-MM"
                    if isinstance(time_key, str):
                        demand[(user_stcd_from_data, time_key)] = demand_value
        else:
            # 如果demand_data是字典格式（旧格式兼容）
            for user_stcd, user_demand in self.demand_data.items():
                # 从配置中获取用户名称
                user_stcd_final = None
                water_users_config = self.config.get('water_users', [])
                if isinstance(water_users_config, dict):
                    water_users_list = water_users_config.values()
                else:
                    water_users_list = water_users_config
                
                for config_user in water_users_list:
                    if config_user.get('user_stcd') == user_stcd:
                        user_stcd_final = config_user.get('user_stcd', user_stcd)
                        break
                
                if not user_stcd_final:
                    user_stcd_final = user_stcd
                    logger.warning(f"未找到用水户 {user_stcd} 的配置信息，使用原名称")
                
                # 处理需水数据
                for time_key, demand_value in user_demand.items():
                    # 直接使用time_key作为时间标识，格式为 "YYYY-MM"
                    if isinstance(time_key, str):
                        demand[(user_stcd_final, time_key)] = demand_value
        
        return demand
    
    def _process_supply_relationship(self) -> Dict[Tuple[str, str], int]:
        """
        处理供给关系
        
        Returns:
            字典，key为(供水水库名称, 用水户名称)的元组，value为0或1
        """
        supply_relationship = {}
        
        # 遍历供给关系配置
        supply_relationships_config = self.config.get('supply_relationships', [])
        if isinstance(supply_relationships_config, dict):
            relationships_list = supply_relationships_config.values()
        else:
            relationships_list = supply_relationships_config
        
        for relationship in relationships_list:
            reservoir_stcd = relationship['reservoir_stcd']
            user_stcd = relationship['user_stcd']
            
            # 设置供给关系为1（表示供水）
            supply_relationship[(reservoir_stcd, user_stcd)] = 1
        
        # 为所有可能的组合设置默认值0
        reservoirs_config = self.config.get('reservoirs', [])
        if isinstance(reservoirs_config, dict):
            all_reservoirs = [res['stcd'] for res in reservoirs_config.values()]
        else:
            all_reservoirs = [res['stcd'] for res in reservoirs_config]
        
        water_users_config = self.config.get('water_users', [])
        if isinstance(water_users_config, dict):
            all_users = [user['user_stcd'] for user in water_users_config.values()]
        else:
            all_users = [user['user_stcd'] for user in water_users_config]
        
        for reservoir_stcd in all_reservoirs:
            for user_stcd in all_users:
                if (reservoir_stcd, user_stcd) not in supply_relationship:
                    supply_relationship[(reservoir_stcd, user_stcd)] = 0
        
        return supply_relationship

    def _process_reservoir_supply(self) -> Dict[str, List[str]]:
        """
        处理每个水库与其对应用水户的关系（按水库聚合）
        
        Returns:
            字典，key为reservoir_stnm（水库名称），value为该水库供水的所有water_user（用水户标识）列表
        """
        reservoir_to_users: Dict[str, List[str]] = {}

        relationships = self.config.get('supply_relationships', {})
        iterable = relationships.values() if isinstance(relationships, dict) else relationships

        for rel in iterable:
            reservoir_stcd = rel.get('reservoir_stcd')
            user_stcd = rel.get('user_stcd')
            if not reservoir_stcd or not user_stcd:
                continue
            users_list = reservoir_to_users.setdefault(reservoir_stcd, [])
            if user_stcd not in users_list:
                users_list.append(user_stcd)

        return reservoir_to_users
    
    def _process_storage_init(self) -> Dict[Tuple[str, str], float]:
        """
        处理初始库容
        
        Returns:
            字典，key为(水库名称, 时间)的元组，value为初始库容
        """
        storage_init = {}
        
        # 获取起始时间点
        start_time = self.config['general_parameters'].get('start_time')
        if start_time:
            start_year = int(start_time.split('-')[0])
            start_month = int(start_time.split('-')[1])
            start_time_key = f"{start_year}-{start_month:02d}"
        else:
            # 如果配置中没有起始时间，尝试从数据中获取第一个时间点
            time_sets = self._process_time_sets()
            if time_sets:
                start_time_key = time_sets[0]
            else:
                raise ValueError("无法确定起始时间：配置文件中缺少 start_time，且数据文件中也没有找到时间信息")
        
        # 从obj_cons_data中读取初始库容数据
        initial_storage_config = self.obj_cons_data.get('constrains', {}).get('initial_storage', [])
        
        # 创建水库名称到初始库容的映射
        storage_map = {}
        for storage_info in initial_storage_config:
            reservoir_stcd = storage_info['stcd']
            init_storage = storage_info['init_storage']
            storage_map[reservoir_stcd] = init_storage
        
        # 遍历所有水库，设置初始库容
        reservoirs_config = self.config.get('reservoirs', [])
        if isinstance(reservoirs_config, dict):
            reservoirs_list = reservoirs_config.values()
        else:
            reservoirs_list = reservoirs_config
        
        for reservoir_info in reservoirs_list:
            reservoir_stcd = reservoir_info['stcd']
            
            # 从obj_cons_data中获取初始库容，如果没有则使用默认值0.0
            initial_storage = storage_map.get(reservoir_stcd, 0.0)
            
            # 设置初始时间点的库容
            storage_init[(reservoir_stcd, start_time_key)] = initial_storage
        
        return storage_init
    
    def _process_end_storage_flood_season_cons(self) -> Dict[str, float]:
        """
        处理汛末库容约束（end_storage_flood_season）
        仅对在 obj_cons.json 中明确配置的水库生成参数；
        未配置的水库不加入到返回字典中（不设默认值）。
        
        该约束要求水库在汛末（10月）的库容等于目标值。

        期望 obj_cons.json 结构示例：
        "constrains": {
            "end_storage_flood_season": [
                {"stnm": "大伙房水库", "target_storage": 1800},
                ...
            ]
        }

        Returns:
            字典，key为水库名称，value为汛末目标库容
        """
        end_storage_flood_season_cons: Dict[str, float] = {}

        # 从 obj_cons_data 中获取 end_storage_flood_season 数据
        if hasattr(self, 'obj_cons_data') and self.obj_cons_data and 'constrains' in self.obj_cons_data:
            end_storage_data = self.obj_cons_data['constrains'].get('end_storage_flood_season', [])
            for reservoir_data in end_storage_data:
                reservoir_stcd = reservoir_data.get('stcd')
                target_storage = reservoir_data.get('target_storage')
                if reservoir_stcd and target_storage is not None:
                    end_storage_flood_season_cons[reservoir_stcd] = target_storage

        return end_storage_flood_season_cons
    
    def _process_end_storage_dispatch_term_cons(self) -> Dict[str, float]:
        """
        处理调度期末库容约束（end_storage_dispatch_term）
        仅对在 obj_cons.json 中明确配置的水库生成参数；
        未配置的水库不加入到返回字典中（不设默认值）。
        
        该约束要求水库在调度期末的库容等于目标值。

        期望 obj_cons.json 结构示例：
        "constrains": {
            "end_storage_dispatch_term": [
                {"stnm": "大伙房水库", "target_storage": 1500},
                ...
            ]
        }

        Returns:
            字典，key为水库名称，value为调度期末目标库容
        """
        end_storage_dispatch_term_cons: Dict[str, float] = {}

        # 从 obj_cons_data 中获取 end_storage_dispatch_term 数据
        if hasattr(self, 'obj_cons_data') and self.obj_cons_data and 'constrains' in self.obj_cons_data:
            end_storage_data = self.obj_cons_data['constrains'].get('end_storage_dispatch_term', [])
            for reservoir_data in end_storage_data:
                reservoir_stcd = reservoir_data.get('stcd')
                target_storage = reservoir_data.get('target_storage')
                if reservoir_stcd and target_storage is not None:
                    end_storage_dispatch_term_cons[reservoir_stcd] = target_storage

        return end_storage_dispatch_term_cons
    
    def _process_storage_up(self) -> Dict[Tuple[str, str], float]:
        """
        处理库容上限
        优先从obj_cons_data中读取，如果没有则使用原有的计算方式
        
        Returns:
            字典，key为(水库名称, 时间)的元组，value为库容上限
        """
        storage_up = {}
        
        # 获取时间序列
        time_sets = self._process_time_sets()
        
        # 创建扩展的时间序列（包含额外的时间点）
        extended_time_series = time_sets.copy()
        
        # 在末尾添加一个额外的时间点
        if extended_time_series:
            last_time = extended_time_series[-1]  # 格式: "YYYY-MM"
            year, month = map(int, last_time.split('-'))
            
            # 计算下一个月
            if month == 12:
                next_year = year + 1
                next_month = 1
            else:
                next_year = year
                next_month = month + 1
            
            next_time = f"{next_year:04d}-{next_month:02d}"
            extended_time_series.append(next_time)
        
        # 首先尝试从obj_cons_data中获取storage_up数据
        obj_cons_storage_up = {}
        if hasattr(self, 'obj_cons_data') and self.obj_cons_data and 'constrains' in self.obj_cons_data:
            storage_up_data = self.obj_cons_data['constrains'].get('storage_up', [])
            for reservoir_data in storage_up_data:
                reservoir_stcd = reservoir_data.get('stcd')
                storage_data = reservoir_data.get('storage_data', {})
                if reservoir_stcd and storage_data:
                    obj_cons_storage_up[reservoir_stcd] = storage_data
        
        # 月份名称到数字的映射
        month_name_to_number = {
            'january': 1, 'february': 2, 'march': 3, 'april': 4,
            'may': 5, 'june': 6, 'july': 7, 'august': 8,
            'september': 9, 'october': 10, 'november': 11, 'december': 12
        }
        
        # 遍历所有水库
        reservoirs_config = self.config.get('reservoirs', [])
        if isinstance(reservoirs_config, dict):
            reservoirs_list = reservoirs_config.values()
        else:
            reservoirs_list = reservoirs_config
        
        for reservoir_info in reservoirs_list:
            reservoir_stcd = reservoir_info['stcd']
            
            # 为每个时间点设置库容上限
            for time_key in extended_time_series:
                # 优先从obj_cons_data中获取数据
                if reservoir_stcd in obj_cons_storage_up and time_key in obj_cons_storage_up[reservoir_stcd]:
                    storage_up[(reservoir_stcd, time_key)] = obj_cons_storage_up[reservoir_stcd][time_key]
                else:
                    # 回退到原有的计算方式
                    normal_storage = reservoir_info['characteristics']['normal_storage']
                    flood_season_storage = reservoir_info.get('flood_season_storage', {})
                    
                    year, month = map(int, time_key.split('-'))
                    
                    # 检查当前月份是否在该水库的汛期库容配置中
                    flood_storage = None
                    for month_name, storage in flood_season_storage.items():
                        if month_name_to_number.get(month_name.lower()) == month:
                            flood_storage = storage
                            break
                    
                    if flood_storage is not None:
                        # 使用汛期库容
                        storage_up[(reservoir_stcd, time_key)] = flood_storage
                    else:
                        # 使用正常库容
                        storage_up[(reservoir_stcd, time_key)] = normal_storage
        
        return storage_up
    
    def _process_storage_down(self) -> Dict[Tuple[str, str], float]:
        """
        处理库容下限
        优先从obj_cons_data中读取，如果没有则使用原有的计算方式
        
        Returns:
            字典，key为(水库名称, 时间)的元组，value为库容下限
        """
        storage_down = {}
        
        # 获取时间序列
        time_sets = self._process_time_sets()
        
        # 创建扩展的时间序列（包含额外的时间点）
        extended_time_series = time_sets.copy()
        
        # 在末尾添加一个额外的时间点
        if extended_time_series:
            last_time = extended_time_series[-1]  # 格式: "YYYY-MM"
            year, month = map(int, last_time.split('-'))
            
            # 计算下一个月
            if month == 12:
                next_year = year + 1
                next_month = 1
            else:
                next_year = year
                next_month = month + 1
            
            next_time = f"{next_year:04d}-{next_month:02d}"
            extended_time_series.append(next_time)
        
        # 首先尝试从obj_cons_data中获取storage_down数据
        obj_cons_storage_down = {}
        if hasattr(self, 'obj_cons_data') and self.obj_cons_data and 'constrains' in self.obj_cons_data:
            storage_down_data = self.obj_cons_data['constrains'].get('storage_down', [])
            for reservoir_data in storage_down_data:
                reservoir_stcd = reservoir_data.get('stcd')
                storage_data = reservoir_data.get('storage_data', {})
                if reservoir_stcd and storage_data:
                    obj_cons_storage_down[reservoir_stcd] = storage_data
        
        # 遍历所有水库
        reservoirs_config = self.config.get('reservoirs', [])
        if isinstance(reservoirs_config, dict):
            reservoirs_list = reservoirs_config.values()
        else:
            reservoirs_list = reservoirs_config
        
        for reservoir_info in reservoirs_list:
            reservoir_stcd = reservoir_info['stcd']
            
            # 为每个时间点设置库容下限
            for time_key in extended_time_series:
                # 优先从obj_cons_data中获取数据
                if reservoir_stcd in obj_cons_storage_down and time_key in obj_cons_storage_down[reservoir_stcd]:
                    storage_down[(reservoir_stcd, time_key)] = obj_cons_storage_down[reservoir_stcd][time_key]
                else:
                    # 回退到原有的计算方式
                    dead_storage = reservoir_info['characteristics']['dead_storage']
                    storage_down[(reservoir_stcd, time_key)] = dead_storage
        
        return storage_down

    def _process_spillflow_up(self) -> Dict[Tuple[str, str], float]:
        """
        处理弃水上限约束（spillflow_up）
        仅对在 obj_cons.json 中明确配置的水库与月份生成参数；
        未配置的水库不加入到返回字典中（不设默认值）。

        Returns:
            字典，key为(水库名称, 时间)的元组，value为弃水约束值
        """
        spillflow_up: Dict[Tuple[str, str], float] = {}

        # 从 obj_cons_data 中获取 spillflow_up 数据（形式同 storage_down，使用 storage_data 月序列）
        obj_cons_spillflow_up: Dict[str, Dict[str, float]] = {}
        if hasattr(self, 'obj_cons_data') and self.obj_cons_data and 'constrains' in self.obj_cons_data:
            spillflow_up_data = self.obj_cons_data['constrains'].get('spillflow_up', [])
            for reservoir_data in spillflow_up_data:
                reservoir_stcd = reservoir_data.get('stcd')
                spill_data = reservoir_data.get('spill_data', {})
                if reservoir_stcd and spill_data:
                    obj_cons_spillflow_up[reservoir_stcd] = spill_data

        # 仅为已配置的水库与其提供的月份生成键值
        for reservoir_stcd, spill_data in obj_cons_spillflow_up.items():
            for time_key, val in spill_data.items():
                spillflow_up[(reservoir_stcd, time_key)] = val

        return spillflow_up

    def _process_spillflow_down(self) -> Dict[Tuple[str, str], float]:
        """
        处理弃水下限约束（spillflow_down）
        仅对在 obj_cons.json 中明确配置的水库与月份生成参数；
        未配置的水库不加入到返回字典中（不设默认值）。

        Returns:
            字典，key为(水库名称, 时间)的元组，value为弃水下限约束值
        """
        spillflow_down: Dict[Tuple[str, str], float] = {}

        # 从 obj_cons_data 中获取 spillflow_down 数据
        obj_cons_spillflow_down: Dict[str, Dict[str, float]] = {}
        if hasattr(self, 'obj_cons_data') and self.obj_cons_data and 'constrains' in self.obj_cons_data:
            spillflow_down_data = self.obj_cons_data['constrains'].get('spillflow_down', [])
            for reservoir_data in spillflow_down_data:
                reservoir_stcd = reservoir_data.get('stcd')
                spill_data = reservoir_data.get('spill_data', {})
                if reservoir_stcd and spill_data:
                    obj_cons_spillflow_down[reservoir_stcd] = spill_data

        # 仅为已配置的水库与其提供的月份生成键值
        for reservoir_stcd, spill_data in obj_cons_spillflow_down.items():
            for time_key, val in spill_data.items():
                spillflow_down[(reservoir_stcd, time_key)] = val

        return spillflow_down

    def _process_import_export_up(self) -> Dict[Tuple[str, str, str], float]:
        """
        import_export_up
        仅对在 obj_cons.json 中明确配置的 (source_stcd, target_reservoir_stcd, YYYY-MM) 生成参数；
        若 json 未配置，则不返回任何键（与 _process_spillflow_up 一致）。

        期望 obj_cons.json 结构示例：
        "constrains": {
            "import_export_up": [
                {"source_stnm": "桓仁", "target_reservoir_stnm": "大伙房水库", "trans_data": {"2024-01": 50.0, ...}},
                ...
            ]
        }

        Returns:
            字典，key为(source_stnm, target_reservoir_stnm, 时间)的元组，value为上限
        """
        import_export_up: Dict[Tuple[str, str, str], float] = {}

        # 从 obj_cons_data 中获取 import_export_up 配置
        if hasattr(self, 'obj_cons_data') and self.obj_cons_data and 'constrains' in self.obj_cons_data:
            import_export_up_list = self.obj_cons_data['constrains'].get('import_export_up', [])
            for entry in import_export_up_list:
                source = entry.get('source_stcd')
                target = entry.get('target_reservoir_stcd')
                trans_data = entry.get('trans_data', {})
                if source and target and trans_data:
                    for time_key, val in trans_data.items():
                        if isinstance(time_key, str):
                            import_export_up[(source, target, time_key)] = val

        return import_export_up
    
    def _process_import_export_down(self) -> Dict[Tuple[str, str, str], float]:
        """
        处理外部水源引入下限约束（import_export_down）
        仅对在 obj_cons.json 中明确配置的 (source_stnm, target_reservoir_stnm, YYYY-MM) 生成参数；
        若 json 未配置，则不返回任何键（与 _process_import_export_up 一致）。

        期望 obj_cons.json 结构示例：
        "constrains": {
            "import_export_down": [
                {"source_stnm": "桓仁", "target_reservoir_stnm": "大伙房水库", "trans_data": {"2024-01": 0.0, ...}},
                ...
            ]
        }

        Returns:
            字典，key为(source_stnm, target_reservoir_stnm, 时间)的元组，value为下限
        """
        import_export_down: Dict[Tuple[str, str, str], float] = {}

        # 从 obj_cons_data 中获取 import_export_down 配置
        if hasattr(self, 'obj_cons_data') and self.obj_cons_data and 'constrains' in self.obj_cons_data:
            import_export_down_list = self.obj_cons_data['constrains'].get('import_export_down', [])
            for entry in import_export_down_list:
                source = entry.get('source_stcd')
                target = entry.get('target_reservoir_stcd')
                trans_data = entry.get('trans_data', {})
                if source and target and trans_data:
                    for time_key, val in trans_data.items():
                        if isinstance(time_key, str):
                            import_export_down[(source, target, time_key)] = val

        return import_export_down
    
    def _process_supply_relationship_type(self) -> Dict[Tuple[str, str], str]:
        """
        处理供水关系类型
        
        Returns:
            字典，key为(供水水库名称, 用水户名称)的元组，value为取水类型
        """
        supply_relationship_type = {}
        
        # 创建用水户名称到取水方式的映射
        user_intake_methods = {}
        water_users_config = self.config.get('water_users', [])
        if isinstance(water_users_config, dict):
            water_users_list = water_users_config.values()
        else:
            water_users_list = water_users_config
        
        for user_info in water_users_list:
            user_stcd = user_info['user_stcd']
            intake_method = user_info.get('intake_method', 'downstream_intake')  # 默认为下游取水
            user_intake_methods[user_stcd] = intake_method
        
        # 遍历供给关系配置
        supply_relationships_config = self.config.get('supply_relationships', [])
        if isinstance(supply_relationships_config, dict):
            relationships_list = supply_relationships_config.values()
        else:
            relationships_list = supply_relationships_config
        
        for relationship in relationships_list:
            reservoir_stcd = relationship['reservoir_stcd']
            user_stcd = relationship['user_stcd']
            
            # 从config的water_users中获取intake_method
            supply_type = user_intake_methods.get(user_stcd, 'downstream_intake')  # 默认为下游取水
            
            supply_relationship_type[(reservoir_stcd, user_stcd)] = supply_type
        
        return supply_relationship_type
    
    def _process_static(self) -> Dict[str, Any]:
        """
        处理静态信息
        
        Returns:
            水库静态信息字典
        """
        reservoirs = self.config.get('reservoirs', {})
        # 将索引转换为按 stcd（水库编码）索引的字典
        static_by_stcd: Dict[str, Any] = {}
        if isinstance(reservoirs, dict):
            for _, info in reservoirs.items():
                stcd = info.get('stcd')
                if stcd:
                    static_by_stcd[stcd] = info
        elif isinstance(reservoirs, list):
            for info in reservoirs:
                stcd = info.get('stcd')
                if stcd:
                    static_by_stcd[stcd] = info
        else:
            # 非预期结构，直接返回原值
            return reservoirs
        return static_by_stcd
    
    def _process_time_sets(self) -> List[str]:
        """
        处理时间序列索引，返回连续的年-月时间序列
        
        Returns:
            时间序列列表，格式为["2025-05", "2025-06", ...]
        """
        time_periods = self._extract_time_periods()
        
        if not time_periods:
            # 如果没有找到时间数据，使用配置中的起始时间
            start_time = self.config['general_parameters'].get('start_time')
            if not start_time:
                raise ValueError("无法生成时间序列：配置文件中缺少 start_time，且数据文件中也没有找到时间信息")
            
            start_year = int(start_time.split('-')[0])
            start_month = int(start_time.split('-')[1])
            # 默认生成12个月的数据
            time_sets = []
            for i in range(12):
                month = start_month + i
                year = start_year
                if month > 12:
                    month = month - 12
                    year = year + 1
                time_sets.append(f"{year:04d}-{month:02d}")
            return time_sets
        
        # 按时间顺序排序
        time_periods.sort()
        
        # 转换为字符串格式
        time_sets = [f"{year:04d}-{month:02d}" for year, month in time_periods]
        
        return time_sets
    
    def _extract_time_periods(self) -> List[Tuple[int, int]]:
        """
        从数据中提取时间周期（年份，月份）元组列表
        
        Returns:
            时间周期列表，格式为[(年份, 月份), ...]
        """
        time_periods = []
        
        # 从径流数据中提取时间周期
        if 'reservoirs' in self.runoff_data:
            # 新格式：从reservoirs列表中提取
            for reservoir in self.runoff_data['reservoirs']:
                runoff_data = reservoir.get('runoff_data', {})
                for time_key in runoff_data.keys():
                    if isinstance(time_key, str) and '-' in time_key:
                        parts = time_key.split('-')
                        if len(parts) >= 2:
                            year = int(parts[0])
                            month = int(parts[1])
                            time_periods.append((year, month))
                break  # 只需要从第一个水库获取时间信息
        else:
            # 旧格式：从字典中提取
            for reservoir_runoff in self.runoff_data.values():
                for time_key in reservoir_runoff.keys():
                    if isinstance(time_key, str) and '-' in time_key:
                        parts = time_key.split('-')
                        if len(parts) >= 2:
                            year = int(parts[0])
                            month = int(parts[1])
                            time_periods.append((year, month))
                break  # 只需要从第一个水库获取时间信息
        
        # 去重并排序
        time_periods = sorted(list(set(time_periods)))
        
        return time_periods
    
    def _process_engineering_hydro(self) -> List[str]:
        """
        处理水库工程信息
        
        Returns:
            所有水库名称列表
        """
        reservoirs = self.config.get('reservoirs', [])
        # 兼容字典和列表两种格式
        if isinstance(reservoirs, dict):
            return [res['stcd'] for res in reservoirs.values()]
        else:
            return [res['stcd'] for res in reservoirs]
    
    def _process_engineering_water_user(self) -> List[str]:
        """
        处理用水户信息
        
        Returns:
            所有用水户名称列表
        """
        water_users = self.config.get('water_users', [])
        # 兼容字典和列表两种格式
        if isinstance(water_users, dict):
            return [user['user_stcd'] for user in water_users.values()]
        else:
            return [user['user_stcd'] for user in water_users]
    
    def _process_external_water_sources(self) -> Dict[str, Dict[str, Any]]:
        """
        处理外部引入水信息
        
        Returns:
            字典，key为source_stcd（外部水源名称），value为该水源的完整信息
        """
        external_sources: Dict[str, Dict[str, Any]] = {}
        sources = self.config.get('external_water_sources', {})
        # 兼容字典或列表结构
        if isinstance(sources, dict):
            iterable = sources.values()
        else:
            iterable = sources
        for source in iterable:
            # 以名称为键，若缺少名称则回退到编码
            key = source.get('source_stcd')
            if not key:
                continue
            # 保留配置中的全部字段
            external_sources[key] = dict(source)
        return external_sources
    
    def _process_imported_reservoir(self) -> List[str]:
        """
        处理外部水源的目标水库信息
        
        Returns:
            所有外部水源目标水库名称的列表（去重）
        """
        imported_reservoirs = []
        sources = self.config.get('external_water_sources', [])
        
        # 兼容字典或列表结构
        if isinstance(sources, dict):
            iterable = sources.values()
        else:
            iterable = sources
            
        for source in iterable:
            target_reservoir = source.get('target_reservoir_stcd')
            if target_reservoir and target_reservoir not in imported_reservoirs:
                imported_reservoirs.append(target_reservoir)
        
        return imported_reservoirs
    
    def _process_import_relationship_source(self) -> Dict[str, List[str]]:
        """
        处理水源与被引入水库的关系（按水源聚合）
        
        Returns:
            字典，key为source_stnm（水源名称），value为该水源引入的所有target_reservoir_stnm（目标水库名称）列表。
            参数名：import_ralation_source
        """
        import_ralation_source: Dict[str, List[str]] = {}
        sources = self.config.get('external_water_sources', [])
    
        # 兼容字典或列表结构
        if isinstance(sources, dict):
            iterable = sources.values()
        else:
            iterable = sources
    
        for source in iterable:
            source_stcd = source.get('source_stcd')
            target_reservoir = source.get('target_reservoir_stcd')
    
            if source_stcd and target_reservoir:
                reservoirs_list = import_ralation_source.setdefault(source_stcd, [])
                if target_reservoir not in reservoirs_list:
                    reservoirs_list.append(target_reservoir)
    
        return import_ralation_source
    
    def _process_import_relationship(self) -> Dict[str, List[str]]:
        """
        处理外部水源与目标水库的关系
        
        Returns:
            字典，key为target_reservoir_stnm（目标水库名称），value为给该水库供水的所有source_stnm（外部水源名称）列表
        """
        import_relationship: Dict[str, List[str]] = {}
        sources = self.config.get('external_water_sources', [])
        
        # 兼容字典或列表结构
        if isinstance(sources, dict):
            iterable = sources.values()
        else:
            iterable = sources
            
        for source in iterable:
            source_stcd = source.get('source_stcd')
            target_reservoir = source.get('target_reservoir_stcd')
            
            if source_stcd and target_reservoir:
                sources_list = import_relationship.setdefault(target_reservoir, [])
                if source_stcd not in sources_list:
                    sources_list.append(source_stcd)
        
        return import_relationship
    
    def _process_loss_rate(self) -> Dict[Tuple[str, str], float]:
        """
        处理损失率
        
        Returns:
            字典，key为(供水水库名称, 用水户名称)的元组，value为损失率
        """
        loss_rate = {}
        
        # 遍历供给关系配置
        supply_relationships_config = self.config.get('supply_relationships', [])
        if isinstance(supply_relationships_config, dict):
            relationships_list = supply_relationships_config.values()
        else:
            relationships_list = supply_relationships_config
        
        for relationship in relationships_list:
            reservoir_stcd = relationship['reservoir_stcd']
            user_stcd = relationship['user_stcd']
            rate = relationship.get('loss_rate', 0.0)
            
            loss_rate[(reservoir_stcd, user_stcd)] = rate
        
        return loss_rate
    
    def _process_trans_capacity(self) -> Dict[Tuple[str, str], float]:
        """
        处理输水能力
        
        Returns:
            字典，key为(供水水库名称, 用水户名称)的元组，value为输水能力
        """
        trans_capacity = {}
        
        # 遍历供给关系配置
        supply_relationships_config = self.config.get('supply_relationships', [])
        if isinstance(supply_relationships_config, dict):
            relationships_list = supply_relationships_config.values()
        else:
            relationships_list = supply_relationships_config
        
        for relationship in relationships_list:
            reservoir_stcd = relationship['reservoir_stcd']
            user_stcd = relationship['user_stcd']
            capacity = relationship.get('trans_capacity', 0.0)
            
            trans_capacity[(reservoir_stcd, user_stcd)] = capacity
        
        return trans_capacity
    
    def _process_reservoir_connections(self) -> Dict[str, List[str]]:
        """
        处理水库间连接关系
        
        Returns:
            字典，key为下游水库名称（stcd），value为上游水库名称（stcd）的列表
        """
        reservoir_connections: Dict[str, List[str]] = {}

        # 遍历水库间连接配置
        connections = self.config.get('inter_reservoir_connections', {})

        # 兼容字典或列表结构
        iterable = connections.values() if isinstance(connections, dict) else connections

        for connection in iterable:
            # 优先使用名称字段，若缺失则回退到编码字段
            from_reservoir = connection.get('from_reservoir_stcd') or connection.get('from_reservoir')
            to_reservoir = connection.get('to_reservoir_stcd') or connection.get('to_reservoir')

            if not to_reservoir or not from_reservoir:
                # 跳过缺少必要信息的连接
                continue

            # 将多个上游水库聚合到同一下游水库的列表中，并去重
            upstream_list = reservoir_connections.setdefault(to_reservoir, [])
            if from_reservoir not in upstream_list:
                upstream_list.append(from_reservoir)

        return reservoir_connections


def process_para_data(para: Dict[str, Any]) -> Dict[str, Any]:
    """
    处理para参数数据的便捷函数
    
    Args:
        para: 从set_up.py获得的参数字典
        
    Returns:
        处理后的数据字典
    """
    processor = DataProcessor(para)
    return processor.process_data()


if __name__ == "__main__":
    # 测试代码
    import sys
    import os
    
    # 添加项目根目录到路径
    project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    sys.path.insert(0, project_root)
    
    from optmodel.set_up import load_water_supply_model
    
    try:
        # 加载配置和数据
        config_path = os.path.join(project_root, "config.json")
        para = load_water_supply_model(config_path)
        
        # 处理数据
        param = process_para_data(para)

        logger.info("\n=== 数据处理完成 ===")
        # logger.info(f"径流数据条目数: {len(param['runoff'])}")
        # logger.info(f"需水数据条目数: {len(param['demand'])}")
        # logger.info(f"供给关系条目数: {len(param['supply_relationship'])}")
        # logger.info(f"损失率条目数: {len(param['loss_rate'])}")
        # logger.info(f"输水能力条目数: {len(param['trans_capacity'])}")
        # logger.info(f"水库连接关系条目数: {len(param['reservoir_connections'])}")
        # logger.info(f"水库数量: {len(param['engineering_hydro'])}")
        # logger.info(f"用水户数量: {len(param['engineering_water_user'])}")
        # logger.info(f"外部水源数量: {len(param['external_water_sources'])}")
        
    except Exception as e:
        logger.error(f"数据处理失败: {e}")